# @file    : test3
# @time    : 2025/4/18
# @author  : yongpeng.yao
# @desc    :
import os
from typing import List, Optional, Union

import dashscope
import pandas as pd
from dashscope import Generation
from dotenv import load_dotenv

load_dotenv()
# 设置通义大模型API密钥
dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")


def read_excel_to_context(file_path: str,
                          sheet_name: Optional[Union[str, int]] = None,
                          text_columns: List[str] = None):
    """
    读取Excel文件并构造上下文字符串

    :param file_path: Excel文件路径
    :param sheet_name: 工作表名称，默认为第一个工作表
    :param text_columns: 需要包含的文本列名列表，如果为None则包含所有列
    :return: 构造好的上下文字符串
    """
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path, sheet_name=sheet_name, engine='openpyxl')
        print(f"成功读取Excel文件，共{len(df)}行数据")

        # 如果没有指定列，则使用所有列
        if text_columns is None:
            text_columns = df.columns.tolist()

        # 构造上下文
        context_parts = []
        for _, row in df.iterrows():
            row_str = " | ".join(f"{col}: {row[col]}" for col in text_columns if pd.notna(row[col]))
            context_parts.append(row_str)

        full_context = "\n".join(context_parts)
        return full_context

    except Exception as e:
        print(f"处理Excel文件时出错: {e}")
        return None


def analyze_with_tongyi(context: str,
                        query: str,
                        model: str = "qwen-max"):
    """
    使用通义大模型分析上下文

    :param context: 上下文字符串
    :param query: 分析查询或指令
    :param model: 使用的模型名称，默认为qwen-max
    :return: 模型响应字典
    """
    try:
        # 构造提示词
        prompt = f"""请根据以下上下文信息进行分析总结：

        上下文数据：
        {context}

        分析任务：{query}

        请提供结构化的分析结果，包括关键发现、趋势分析和建议。"""

        # 调用通义大模型
        response = Generation.call(
            model=model,
            prompt=prompt,
            max_tokens=2000,
            temperature=0.7,
            top_p=0.9
        )

        return response

    except Exception as e:
        print(f"调用通义大模型时出错: {e}")
        return None


def main():
    # 配置参数
    excel_path = "./data/translated2.xlsx"  # 替换为你的Excel文件路径
    sheet_name = 0  # 替换为你的工作表名或None使用第一个工作表

    # 1. 读取Excel并构造上下文
    context = read_excel_to_context(excel_path, sheet_name)
    if context is None:
        print("读取Excel文件失败，请检查文件路径和列名")
        return

    print("\n上下文数据预览：")
    print(context[:500] + "...")  # 打印前500个字符作为预览

    # 2. 交互式分析
    while True:
        print("\n请输入分析指令或问题(输入'exit'退出)：")
        user_query = input().strip()

        if user_query.lower() == 'exit':
            break

        if not user_query:
            print("指令不能为空")
            continue

        # 3. 使用通义大模型分析
        print("\n分析结果：")
        result = analyze_with_tongyi(context, user_query)

        if result and 'output' in result:
            print(result['output']['text'])
        else:
            print("未能获取有效分析结果")


if __name__ == "__main__":
    main()
